from __future__ import annotations
import inspect
from queue import Queue
import time
from weakref import WeakValueDictionary, WeakSet
from pathlib import Path
from typing import Literal, Optional, Union, Set
from logging import Logger, Formatter, StreamHandler, FileHandler
from logging.handlers import RotatingFileHandler, QueueHandler, QueueListener
from threading import Event
import logging
import sys
import gc


# %% Get Manager Instance
_logger_manager = logging.getLoggerClass().manager

# %% File Handlers Manager (Singleton)
# 管理文件句柄, 通过QueueHandler实现异步日志处理
class FileHandlerManager:
    _instance = None
    _handlers: WeakValueDictionary  # name -> filehandler
    _handlers_set: Set  # filehandler
    _queues: Queue  # listener queue
    _listeners: QueueListener
    _queue_handler: QueueHandler

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._handlers = WeakValueDictionary()
            cls._handlers_set = set()
            cls._queues = Queue(-1)
            cls._queue_handler = QueueHandler(cls._queues)
            cls._queue_handler.setFormatter(file_formatter)
            cls._listener = cls.start_listener(cls._queues, cls._handlers_set)
        return cls._instance

    def __contains__(self, query: Union[str, FileHandler]) -> bool:
        if isinstance(query, str):
            return query in self._handlers.keys()
        elif isinstance(query, FileHandler):
            return query in self._handlers.values()

    @staticmethod
    def start_listener(
        queue: Queue,
        handlers_set: WeakSet[FileHandler],
    ) -> QueueListener:
        listener = QueueListener(queue, respect_handler_level=True)
        listener.handlers = handlers_set  # 绑定线程监听对象
        listener.start()  # 自动启动监听线程
        return listener

    def getHandler(
        self,
        handler=FileHandler,
    ) -> QueueHandler:
        key = handler.baseFilename
        if key in self._handlers:
            assert isinstance(handler, type(self._handlers[key])), (
                f"exists handler {key} is not {type(handler)}, got {type(self._handlers[key])}"
            )
            print("[FileHandlerManager] Reusing existing handler:", key)
        else:
            self._handlers_set.add(handler)
            self._handlers[key] = handler
        return self._queue_handler

    def removeHandler(
        self,
        handler: FileHandler,
    ) -> None:
        key = handler.baseFilename
        if key in self._handlers:
            self._handlers_set.discard(handler)
            gc.collect()
        else:
            raise ValueError(f"Handler {key} not found in handlers set")


# %% Formatter
class ColoredFormatter(Formatter):
    COLORS = {
        "WARNING": "\033[33m",  # Yellow
        "ERROR": "\033[31m",  # Red
        "CRITICAL": "\033[31m",  # Red
        "DEFAULT": "\033[32m",  # Green
        "INFO": "\033[34m",  # Blue
    }
    RESET = "\033[0m"
    NAME_WIDTH = 0

    def formatTime(self, record, datefmt=None):
        ct = time.localtime(record.created)
        ms = int((record.created - int(record.created)) * 1e3)
        if datefmt:
            s = time.strftime(datefmt, ct)
            s = f"{s}.{ms:03d}"
        else:
            s = time.strftime("%Y-%m-%d %H:%M:%S", ct)
            s = f"{s}.{ms:03d}"
        return s

    def format(self, record):
        color = self.COLORS.get(record.levelname, self.COLORS["DEFAULT"])
        record.message = record.getMessage().strip()
        record.asctime = self.formatTime(record, self.datefmt)
        self.__class__.NAME_WIDTH = max(self.__class__.NAME_WIDTH, len(record.name))
        s = f"{record.asctime} | {color}{record.levelname:<8}{self.RESET} | {record.name:<{self.__class__.NAME_WIDTH}} | {record.message}"
        return s


class FileFormatter(Formatter):
    NAME_WIDTH = 0

    def formatTime(self, record, datefmt=None):
        ct = time.localtime(record.created)
        ms = int((record.created - int(record.created)) * 1e3)
        if datefmt:
            s = time.strftime(datefmt, ct)
            s = f"{s}.{ms:03d}"
        else:
            s = time.strftime("%Y-%m-%d %H:%M:%S", ct)
            s = f"{s}.{ms:03d}"
        return s

    def format(self, record):
        record.message = record.getMessage().strip()
        record.asctime = self.formatTime(record, self.datefmt)
        self.__class__.NAME_WIDTH = max(self.__class__.NAME_WIDTH, len(record.name))
        s = f"{record.asctime} | {record.levelname:<8} | {record.name:<{self.__class__.NAME_WIDTH}} | {record.message}"
        return s

file_formatter = FileFormatter()

class AsyncLogger(Logger):

    def __call__(self, string: str = "", level: int = logging.DEBUG):
        super().log(level, string)

    def addHandler(self, hdlr):
        if isinstance(hdlr, FileHandler):
            handler_manager = FileHandlerManager()
            queue_hdlr = handler_manager.getHandler(hdlr)
            return super().addHandler(queue_hdlr)
        else:
            return super().addHandler(hdlr)

    def removeHandler(self, hdlr):
        if isinstance(hdlr, FileHandler):
            handler_manager = FileHandlerManager()
            handler_manager.removeHandler(hdlr)
        else:
            return super().removeHandler(hdlr)

    def addFileHandler(
        self,
        log_file: Union[str, Path],
        type: Literal["base", "rotating", "time"] = "rotating",
        level=logging.NOTSET,
        maxBytes=10485760,
        backupCount=10,
    ):
        if isinstance(log_file, str):
            log_file = Path(log_file).expanduser().resolve()
        log_file = str(log_file)
        if type == "base":
            file_handler = FileHandler(log_file)
        elif type == "rotating":
            file_handler = RotatingFileHandler(
                log_file, maxBytes=maxBytes, backupCount=backupCount
            )
            file_handler.setLevel(level)
        elif type == "time":
            raise NotImplementedError("Time-based file handler not implemented yet")
        return self.addHandler(file_handler)

    def error(
        self, msg, *args, exc_info=None, stack_info=False, stacklevel=1, extra=None
    ):
        frame = inspect.currentframe().f_back
        filename = frame.f_code.co_filename
        lineno = frame.f_lineno
        msg = f"{filename}:{lineno} | {msg}"
        return super().error(
            msg,
            *args,
            exc_info=exc_info,
            stack_info=stack_info,
            stacklevel=stacklevel,
            extra=extra,
        )

    def stopwatch(this, timeout_ms: Optional[int] = None):
        class StopwatchContext:
            def __enter__(self, *args, **kwargs):
                self.timeout_ms = timeout_ms
                self.start_time = time.monotonic()
                return self

            def __exit__(self, exc_type, exc_value, traceback):
                frame = inspect.currentframe().f_back
                filename = frame.f_code.co_filename
                lineno = frame.f_lineno
                end_time = time.monotonic()
                duration_ms = int((float(end_time) - float(self.start_time)) * 1000)
                if self.timeout_ms is not None and duration_ms > self.timeout_ms:
                    this(
                        f"{filename}:{lineno} | TAKE {duration_ms} MS > TIMEOUT {self.timeout_ms} MS",
                        logging.WARNING,
                    )
                if self.timeout_ms is None:
                    this(f"{filename}:{lineno} | TAKE {duration_ms} MS", logging.DEBUG)
                return False

        return StopwatchContext()

# %% Global Logger Instance
_logger_manager.setLoggerClass(AsyncLogger) # 更改默认Logger类

# %% functions
_is_initialized = Event()
def get_logger(name, *, level=logging.DEBUG, log_file=None, init_basic=True) -> AsyncLogger:
    logger:AsyncLogger = _logger_manager.getLogger(name)
    logger.setLevel(level)
    # Check if handlers already exist to prevent duplicates
    if init_basic:
        if not _is_initialized.is_set():
            formatter = ColoredFormatter()
            handler = StreamHandler(sys.stdout)
            handler.setFormatter(formatter)
            logging.basicConfig(level=level, handlers=[handler])
            _is_initialized.set()
    else:
        if not logger.handlers:
            formatter = ColoredFormatter()
            handler = StreamHandler(sys.stdout)
            handler.setFormatter(formatter)
            logger.addHandler(handler)
    # File handler for file output
    if log_file:
        logger.addFileHandler(log_file, level=level)
    return logger